"""SQL Agent工具模块 - 将SQL Agent包装成可调用的工具。"""

from typing import Annotated
from langchain_core.tools import tool
from langchain_core.runnables import RunnableConfig
from agents.states import SqlGraphAnnotation
from agents.configurations import SqlConfiguration


@tool
async def sql_query_tool(
    query: Annotated[str, "用户的SQL查询请求或数据库相关问题"]
) -> str:
    """执行SQL查询或回答数据库相关问题。
    
    这个工具可以：
    1. 理解用户的自然语言查询需求
    2. 自动生成相应的SQL查询
    3. 执行查询并返回结果
    4. 处理数据库schema相关问题
    
    Args:
        query: 用户的查询请求，可以是自然语言描述的数据需求
        
    Returns:
        str: 查询结果或相关信息
    """
    try:
        # 延迟导入避免循环依赖
        from agents.graphs import sql_graph
        
        # 创建SQL Agent的初始状态
        initial_state = SqlGraphAnnotation(
            messages=[
                {"role": "user", "content": query}
            ]
        )
        
        # 创建配置
        config = RunnableConfig(
            configurable={
                "model": "deepseek-ai/DeepSeek-V3",
                "system_prompt": "You are a helpful SQL assistant.",
                "database_uri": "sqlite:///Chinook.db"
            }
        )
        
        # 调用SQL Agent子图
        result = await sql_graph.ainvoke(initial_state, config)
        
        # 提取最后的消息作为结果
        if result.messages:
            last_message = result.messages[-1]
            if hasattr(last_message, 'content'):
                return str(last_message.content)
        
        return "SQL查询执行完成，但没有返回结果。"
        
    except Exception as e:
        return f"SQL查询执行出错: {str(e)}"